#! /usr/bin/python3

import rospy
from ocs2_msgs.msg import mpc_observation, mpc_flattened_controller, mpc_state, mpc_input
from std_msgs.msg import Float64
from sensor_msgs.msg import JointState


def policyCallback(msg: mpc_flattened_controller):
    # 猜测 data 0 就是输入的力
    out: Float64 = msg.data[0].data[0]
    pub_to_gazebo.publish(out)


def inputCallback(msg: mpc_input):
    out = msg.value[0]
    pub_to_gazebo.publish(out)


def stateCallback(msg: JointState):
    jointName = [
        'slider_to_cart', 'cart_to_pole1', 'pole1_to_pole2', 'pole2_to_pole3'
    ]
    # index = 0
    # for i in JointState.name:
    #     index = jointName.index(i)
    # 猜测 先第一个关节，再第二个关节
    x = []
    for i in range(len(jointName)):
        index = msg.name.index(jointName[i])
        x.append(msg.position[index])

    for i in range(len(jointName)):
        index = msg.name.index(jointName[i])
        x.append(msg.velocity[index])
    out: mpc_observation = mpc_observation()
    time = msg.header.stamp
    out.time = time.secs + time.nsecs * 1e-9
    out.input = mpc_input([0])
    out.state = mpc_state(x)
    pub_to_ocs2.publish(out)


if __name__ == "__main__":
    rospy.init_node("ros_toocs2_bridge")
    rate = rospy.Rate(500)
    # observation
    pub_to_ocs2 = rospy.Publisher("/cartpole_mpc_observation",
                                  mpc_observation,
                                  queue_size=2)
    sub_from_ocs2 = rospy.Subscriber("/cartpole_mpc_policy",
                                     mpc_flattened_controller, policyCallback)
    sub_input_from_ocs2 = rospy.Subscriber("/cartpole_mpc_input", mpc_input,
                                           inputCallback)

    pub_to_gazebo = rospy.Publisher("/cartpole/cart_effort_controller/command",
                                    Float64,
                                    queue_size=2)
    sub_from_gazebo = rospy.Subscriber("/cartpole/joint_states", JointState,
                                       stateCallback)
    # 循环
    while not rospy.is_shutdown():
        try:
            rate.sleep()
        except rospy.exceptions.ROSInterruptException:
            print("程序退出")
            break
